#include	"omni.h"
#include	"router.h"
#include	"response.h"
#include	"utils/command.h"
#include	"utils/logger.h"
#include	"utils/utils.h"

extern "C" {
#	include	<luajit.h>
#	include	<lualib.h>
#	include	<lauxlib.h>
}

#include	<cstdio>
#include	<cstring>
#include	<map>
#include	<vector>

#if defined(_WIN32)
#	include	<Windows.h>
#	define	YIELD(n)	Sleep(n)
#	pragma	warning(disable:4996)
#else
#	include	<unistd.h>
#	include	<netinet/in.h>
#	define	YIELD(n)	usleep(n)
#endif

typedef std::map<std::string, std::string> Parameter;
typedef std::map<std::string, std::vector<std::string>> ArrayParameter;

struct Context {
	double				nStart;
	MHD_PostProcessor *	pp;
	FILE *				pf;
	std::string			sFile;
	Parameter			mParam;
	ArrayParameter		mArray;
	Parameter			mFile;
};

namespace Handler {
	static std::string GStaticFileDir = "/www/";
	static std::string GUploadFileDir = "./upload/";

	static bool IsResource(const std::string & sUrl) {
		return (sUrl.find(GStaticFileDir) == 0 || sUrl == "/favicon.ico");
	}

	static int ParamIterator(void * p, MHD_ValueKind emKind, const char * pKey, const char * pVal) {
		lua_State * pL = (lua_State *)p;
		size_t nSize = strlen(pKey);

		if (nSize > 2 && pKey[nSize - 2] == '[' && pKey[nSize - 1] == ']') {
			std::string sKey(pKey, nSize - 2);

			lua_getfield(pL, -1, sKey.c_str());
			if (!lua_istable(pL, -1)) {
				lua_pop(pL, 1);
				lua_newtable(pL);
				lua_pushstring(pL, pVal);
				lua_rawseti(pL, -2, 1);
				lua_setfield(pL, -2, sKey.c_str());
			} else {
				int n = (int)lua_objlen(pL, -1);
				lua_pushstring(pL, pVal);
				lua_rawseti(pL, -2, n + 1);
				lua_pop(pL, 1);
			}
		} else {
			lua_pushstring(pL, pVal);
			lua_setfield(pL, -2, pKey);
		}

		return MHD_YES;
	}

	static int PostIterator(void * p, MHD_ValueKind emKind, const char * pKey, const char * pFile, const char * pType, const char * pEncoding, const char * pData, uint64_t nOffset, size_t nSize) {
		if (nSize == 0) return MHD_NO;

		Context * pCtx = (Context *)p;
		if (!pFile) {
			size_t nKey = strlen(pKey);
			if (nKey > 2 && pKey[nKey - 2] == '[' && pKey[nKey - 1] == ']') {
				std::string sKey(pKey, nKey - 2);
				if (pCtx->mArray.find(sKey) == pCtx->mArray.end()) pCtx->mArray[sKey] = std::vector<std::string>();
				pCtx->mArray[sKey].push_back(pData);
			} else {
				if (pCtx->mParam.find(pKey) == pCtx->mParam.end()) pCtx->mParam[pKey] = "";
				pCtx->mParam[pKey].append(pData, nSize);
			}
			return MHD_YES;
		}

		pCtx->mParam[pKey] = pFile;

		if (strlen(pFile) <= 0) {
			GLog.Error("Bad file name to upload!");
			return MHD_NO;
		}

		if (pCtx->sFile != pFile) {
			if (pCtx->pf) fclose(pCtx->pf);

			std::string sTemp = GUploadFileDir + CreateID().substr(0, 16);
			pCtx->mFile[pKey] = sTemp;

			if (!Exists(GUploadFileDir.c_str())) MakeDir(GUploadFileDir.c_str());

			pCtx->sFile = pFile;
			pCtx->pf = fopen(sTemp.c_str(), "wb+");

			if (!pCtx->pf) {
				GLog.Error("Failed to create temp file : %s to store %s", sTemp.c_str(), pFile);
				return MHD_NO;
			}
		}

		if (fwrite(pData, 1, nSize, pCtx->pf) < nSize) {
			GLog.Error("Failed to write data to store %s", pFile);
			return MHD_NO;
		}

		return MHD_YES;
	}

	static int OnReceive(void * p, MHD_Connection * pConn, const char * pUrl, const char * pMethod, const char * pVersion, const char * pData, size_t * pSize, void ** pConnData) {
		double nStart = Tick();
		int nProc = 0;
		std::smatch iMatch;

		if (IsResource(pUrl)) {
			Response iRsp(pConn);
			iRsp.File(pUrl);
			return MHD_YES;
		}

		int nStatus = Router::Get().Match(pMethod, pUrl, nProc, iMatch);
		if (nStatus != MHD_HTTP_OK) {
			Response iRsp(pConn);
			iRsp.Error(nStatus);
			iRsp.Send();
			return MHD_YES;
		}

		if (strcmp(pMethod, "POST") == 0) {
			Context * pCtx = (Context *)(*pConnData);
			if (!pCtx) {
				pCtx = new Context;

				pCtx->nStart	= nStart;
				pCtx->pp		= MHD_create_post_processor(pConn, 65536, &Handler::PostIterator, pCtx);
				pCtx->pf		= NULL;

				*pConnData = pCtx;
				return MHD_YES;
			}

			if (*pSize != 0) {
				MHD_post_process(pCtx->pp, pData, *pSize);
				*pSize = 0;
				return MHD_YES;
			}

			nStart = pCtx->nStart;
		}

		((Omni *)p)->Serve(pConn, pMethod, pUrl, nProc, iMatch, pConnData);
		GLog.Info("/%s\t%.3lfms\t%s", pMethod, Tick() - nStart, pUrl);
		return MHD_YES;
	}

	static void OnComplete(void * p, MHD_Connection * pConn, void ** pConnData, MHD_RequestTerminationCode emCode) {
		Context * pCtx = (Context *)(*pConnData);
		if (pCtx) {
			if (pCtx->pp) MHD_destroy_post_processor(pCtx->pp);
			if (pCtx->pf) fclose(pCtx->pf);
			pCtx->pp = NULL;
			delete pCtx;
		}
		*pConnData = NULL;
	}
}

Omni::Omni() : _pServer(nullptr), _pL(nullptr) {
}

Omni::~Omni() {
	if (_pServer) MHD_stop_daemon(_pServer);
	if (_pL) lua_close(_pL);
}

void Omni::Start(int nArgc, char * pArgv[]) {
	Command iCmd(nArgc, pArgv);

	if (iCmd.Has("-h") || iCmd.Has("help")) {
		printf(
			"omni - Tiny HTTP web server for LUA developer. \n"
			"Author : longshuang@msn.cn.\n"
			"\n[USAGE]\n"
			"-h, help\t\tShow this message.\n"
			"port=[N]\t\tSet listen port.\n"
			"www=[root]\t\t Set static resource root[default is www].\n"
			"upload=[root]\t\t Set static resource root[default is upload].\n"
			"daemon\t\t Run server as daemon application.\n\n");
		return;
	}

	int nPort = iCmd.Has("port") ? atoi(iCmd.Get("port").c_str()) : 80;
	_pServer = MHD_start_daemon(
		MHD_USE_THREAD_PER_CONNECTION | MHD_USE_INTERNAL_POLLING_THREAD | MHD_USE_ERROR_LOG,
		(uint16_t)nPort,
		NULL, NULL,
		&Handler::OnReceive, this,
		MHD_OPTION_CONNECTION_TIMEOUT, (unsigned int)120,
		MHD_OPTION_NOTIFY_COMPLETED, &Handler::OnComplete, NULL,
		MHD_OPTION_END);

	if (!_pServer) {
		GLog.Error("Failed to start HTTP server!");
		return;
	}

	if (iCmd.Has("www")) Handler::GStaticFileDir = "/" + iCmd.Get("www") + "/";
	if (iCmd.Has("upload")) Handler::GUploadFileDir = "./" + iCmd.Get("upload") + "/";
	if (iCmd.Has("daemon")) {
#if defined(_WIN32)
		HWND hWnd = FindWindow("ConsoleWindowClass", NULL);
		if (hWnd) ShowWindow(hWnd, SW_HIDE);
		else GLog.Error("Failed to run application as daemon.");
#else
		if (daemon(1, 0) != 0) GLog.Error("Failed to run application as daemon.");
#endif
	}

	_pL = luaL_newstate();
	luaL_openlibs(_pL);

	extern void RegisterApi_Logger(lua_State *);
	extern void RegisterApi_Utils(lua_State *);
	extern void RegisterApi_Json(lua_State *);
	extern void RegisterApi_Storage(lua_State *);
	extern void RegisterApi_Router(lua_State *);
	extern void RegisterApi_Response(lua_State *);
	RegisterApi_Logger(_pL);
	RegisterApi_Utils(_pL);
	RegisterApi_Json(_pL);
	RegisterApi_Storage(_pL);
	RegisterApi_Router(_pL);
	RegisterApi_Response(_pL);

	int nTop = lua_gettop(_pL);
	lua_getglobal(_pL, "debug");
	lua_getfield(_pL, -1, "traceback");
	lua_remove(_pL, -2);

	extern std::string LUA_TEMPLATE_ENGINE;
	if (0 != luaL_loadstring(_pL, LUA_TEMPLATE_ENGINE.c_str()) || 0 != lua_pcall(_pL, 0, 0, nTop + 1)) {
		GLog.Error("Load template engine : %s", lua_tostring(_pL, -1));
	}

	if (0 != luaL_loadfile(_pL, "scripts/app.lua") || 0 != lua_pcall(_pL, 0, LUA_MULTRET, nTop + 1)) {
		GLog.Error("Load app.lua failed : %s", lua_tostring(_pL, -1));
	}

	lua_settop(_pL, nTop);

	GLog.Info("Server started ...");
	while (true) YIELD(1);
}

void Omni::Serve(MHD_Connection * pConn, const char * pMethod, const char * pUrl, int nProc, std::smatch & rMatch, void ** pConnData) {
	Response iResponse(pConn);
	lua_State * p = lua_newthread(_pL);
	int nTop = lua_gettop(p);

	lua_getglobal(p, "debug");
	lua_getfield(p, -1, "traceback");
	lua_remove(p, -2);
	lua_rawgeti(p, LUA_REGISTRYINDEX, nProc);

	__PrepareRequest(p, pConn, pMethod, pUrl, pConnData);
	__PrepareResponse(p, &iResponse);
	__PrepareExtra(p, rMatch);

	if (lua_pcall(p, 1 + rMatch.size(), 0, nTop + 1) != 0) {
		GLog.Error("Dispatch error : %s\n", lua_tostring(p, -1));
		iResponse.Error(500);
	}

	lua_remove(p, nTop + 1);
	iResponse.Send();
}

void Omni::__PrepareRequest(lua_State * p, MHD_Connection * pConn, const char * pMethod, const char * pUrl, void ** pConnData) {
	Context * pCtx = (Context *)(*pConnData);
	const MHD_ConnectionInfo * pInfo = MHD_get_connection_info(pConn, MHD_CONNECTION_INFO_CLIENT_ADDRESS);
	uint32_t nIP = ((sockaddr_in *)(pInfo->client_addr))->sin_addr.s_addr;
	char sIP[16] = { 0 };
	snprintf(sIP, 16, "%d.%d.%d.%d", (nIP & 0xFF), (nIP & 0xFF00) >> 8, (nIP & 0xFF0000) >> 16, (nIP & 0xFF000000) >> 24);

	lua_newtable(p);
	lua_pushstring(p, pMethod);
	lua_setfield(p, -2, "method");
	lua_pushstring(p, pUrl);
	lua_setfield(p, -2, "url");
	lua_pushstring(p, sIP);
	lua_setfield(p, -2, "remote");
	lua_newtable(p);
	MHD_get_connection_values(pConn, MHD_GET_ARGUMENT_KIND, &Handler::ParamIterator, p);
	lua_setfield(p, -2, "get");
	lua_newtable(p);
	MHD_get_connection_values(pConn, MHD_COOKIE_KIND, &Handler::ParamIterator, p);
	lua_setfield(p, -2, "cookie");
	lua_newtable(p);
	if (pCtx) {
		for (auto & kv : pCtx->mParam) {
			lua_pushstring(p, kv.second.c_str());
			lua_setfield(p, -2, kv.first.c_str());
		}

		for (auto & kv : pCtx->mArray) {
			lua_newtable(p);
			for (size_t i = 1; i <= kv.second.size(); ++i) {
				lua_pushstring(p, kv.second[i - 1].c_str());
				lua_rawseti(p, -2, (int)i);
			}
			lua_setfield(p, -2, kv.first.c_str());
		}
	}
	lua_setfield(p, -2, "post");
	lua_newtable(p);
	if (pCtx) {
		for (auto & kv : pCtx->mFile) {
			lua_pushstring(p, kv.second.c_str());
			lua_setfield(p, -2, kv.first.c_str());
		}
	}
	lua_setfield(p, -2, "file");
}

void Omni::__PrepareResponse(lua_State * p, Response * pResponse) {
	Response ** pRsp = (Response **)(lua_newuserdata(p, sizeof(Response *)));
	(*pRsp) = pResponse;
	luaL_getmetatable(p, "__omni_response");
	lua_setmetatable(p, -2);
}

void Omni::__PrepareExtra(lua_State * p, std::smatch & rMatch) {
	for (size_t i = 1; i < rMatch.size(); ++i) lua_pushstring(p, rMatch[i].str().c_str());
}

int main(int nArgc, char * pArgv[]) {
	Omni iService;
	iService.Start(nArgc, pArgv);
	return 0;
}